//===- Container.h --------------------------------------------------------===//
//===----------------------------------------------------------------------===//
//
// Container descriptor.
//
//===----------------------------------------------------------------------===//

#ifndef INTERFACE_CORE_CONTAINER
#define INTERFACE_CORE_CONTAINER

#include <cstdint>
#include <cassert>
// MemRef descriptor.
// - T represents the type of the elements.
// - N represents the number of dimensions.
// - The storage order is NCHW.
template <typename T, size_t N> class MemRef {
public:
  // Constructor from shape.
  MemRef(T* data, intptr_t size);
  // Constructor from shape.
  MemRef(size_t size, T init = T(0));
  // Get the element at index.
  const T &operator[](size_t index) const;
  T &operator[](size_t index);

protected:
  // Default constructor.
  // This constructor is designed for derived domain-specific constructor.
  MemRef(){};
  // Data.
  // The `aligned` and `allocated` members point to the same address, `aligned`
  // member is responsible for handling data, and `allocated` member is
  // resposible for handling the memory space.
  T *allocated = nullptr;
  T *aligned = nullptr;
  // Offset.
  intptr_t offset = 0;
  // Shape.
  intptr_t sizes[N];
  // Strides.
  intptr_t strides[N];
};

// MemRef Shape Constructor.
// Construct a MemRef object from the data shape and initial value.
// The default initial value is 0.
template <typename T, std::size_t N>
MemRef<T, N>::MemRef(size_t size, T init) {
  this->sizes[0] = size;
  allocated = new T[size];
  aligned = allocated;
  std::fill(aligned, aligned + size, init);
}

template <typename T, std::size_t N>
MemRef<T, N>::MemRef(T* data, intptr_t size) {
  this->sizes[0] = size;
  this->strides[0] = 1;
  allocated = data;
  aligned = allocated;
}

// Get the element at index.
// Return the specific element if the container data size is greater than zero.
// If the data size is negative or zero, which means no space is allocated for
// the container data pointer, this operator does not allow to return the data
// element.
template <typename T, std::size_t N>
const T &MemRef<T, N>::operator[](size_t index) const {
  assert((sizes[0] > 0) && "Invalid container data size.");
  return aligned[index + offset];
}
template <typename T, std::size_t N> T &MemRef<T, N>::operator[](size_t index) {
  assert((sizes[0] > 0) && "Invalid container data size.");
  return aligned[index + offset];
}


#endif
